import os
import re
import tarfile
from urllib.request import urlretrieve
from collections import defaultdict
from math import log
from sklearn.model_selection import train_test_split

def download_and_extract_data(url, extract_to):
    """Download and extract the 20 Newsgroups dataset."""
    if not os.path.exists(extract_to):
        os.makedirs(extract_to)
    tar_path = os.path.join(extract_to, "20news.tar.gz")
    
    # Download dataset
    if not os.path.exists(tar_path):
        print("Downloading dataset...")
        urlretrieve(url, tar_path)
        print("Download complete.")
    
    # Extract dataset
    print("Extracting dataset...")
    with tarfile.open(tar_path, "r:gz") as tar:
        tar.extractall(path=extract_to)
    print("Extraction complete.")

def preprocess_text(text):
    """Preprocess text by lowercasing and removing non-alphanumeric characters."""
    return re.sub(r'[^a-zA-Z0-9\s]', '', text.lower())

def load_data(directory):
    """Load data from 20Newsgroups directory."""
    data, labels = [], []
    label_names = os.listdir(directory)
    label_map = {name: i for i, name in enumerate(label_names)}
    for label_name in label_names:
        label_dir = os.path.join(directory, label_name)
        if os.path.isdir(label_dir):
            for file in os.listdir(label_dir):
                file_path = os.path.join(label_dir, file)
                try:
                    with open(file_path, 'r', encoding='latin-1') as f:
                        data.append(preprocess_text(f.read()))
                        labels.append(label_map[label_name])
                except Exception as e:
                    print(f"Error reading file {file_path}: {e}")
    return data, labels, label_map

def train_naive_bayes(data, labels, num_classes):
    """Train a Naive Bayes classifier."""
    word_counts = [defaultdict(int) for _ in range(num_classes)]
    class_counts = [0] * num_classes
    vocab = set()
    
    for text, label in zip(data, labels):
        class_counts[label] += 1
        words = text.split()
        for word in words:
            word_counts[label][word] += 1
            vocab.add(word)
    
    total_words_per_class = [sum(word_counts[i].values()) for i in range(num_classes)]
    return word_counts, class_counts, total_words_per_class, vocab

def predict_naive_bayes(text, word_counts, class_counts, total_words_per_class, vocab, num_classes):
    """Predict the class of a given text using Naive Bayes."""
    words = text.split()
    log_probs = []
    total_samples = sum(class_counts)
    
    for c in range(num_classes):
        log_prob = log(class_counts[c] / total_samples)
        for word in words:
            word_freq = word_counts[c][word] + 1
            log_prob += log(word_freq / (total_words_per_class[c] + len(vocab)))
        log_probs.append(log_prob)
    
    return log_probs.index(max(log_probs))

def main():
    # Step 1: Download and extract the dataset
    dataset_url = "http://qwone.com/~jason/20Newsgroups/20news-18828.tar.gz"
    data_dir = "./20newsgroups_data"
    download_and_extract_data(dataset_url, data_dir)
    
    # Step 2: Load the dataset
    dataset_path = os.path.join(data_dir, "20news-18828")
    data, labels, label_map = load_data(dataset_path)
    
    # Step 3: Split data into training and testing sets
    train_data, test_data, train_labels, test_labels = train_test_split(data, labels, test_size=0.2, random_state=42)
    
    # Step 4: Train the Naive Bayes classifier
    word_counts, class_counts, total_words_per_class, vocab = train_naive_bayes(train_data, train_labels, len(label_map))
    
    # Step 5: Test the classifier
    correct = 0
    for text, label in zip(test_data, test_labels):
        prediction = predict_naive_bayes(text, word_counts, class_counts, total_words_per_class, vocab, len(label_map))
        if prediction == label:
            correct += 1
    
    accuracy = correct / len(test_labels)
    print(f"Accuracy: {accuracy:.4f}")

if __name__ == "__main__":
    main()
